#!/usr/bin/env node

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import {
  CallToolRequestSchema,
  ErrorCode,
  ListToolsRequestSchema,
  McpError,
} from '@modelcontextprotocol/sdk/types.js';
import mysql from 'mysql2/promise';
import dotenv from 'dotenv';

// 加载环境变量
dotenv.config();

/**
 * MySQL连接池
 */
let pool = null;

/**
 * 初始化MySQL连接池
 * @returns {Promise<mysql.Pool>} MySQL连接池
 */
async function initializeDatabase() {
  try {
    pool = mysql.createPool({
      host: process.env.MYSQL_HOST || 'localhost',
      port: process.env.MYSQL_PORT || 3306,
      user: process.env.MYSQL_USER || 'root',
      password: process.env.MYSQL_PASSWORD || '',
      database: process.env.MYSQL_DATABASE || 'test',
      waitForConnections: true,
      connectionLimit: 10,
      queueLimit: 0,
      acquireTimeout: 60000,
      timeout: 60000,
    });

    // 测试连接
    const connection = await pool.getConnection();
    await connection.ping();
    connection.release();

    console.log('MySQL连接池初始化成功');
    return pool;
  } catch (error) {
    console.error('MySQL连接失败:', error.message);
    throw error;
  }
}

/**
 * 执行SQL查询
 * @param {string} query - SQL查询语句
 * @param {Array} params - 查询参数
 * @returns {Promise<Object>} 查询结果
 */
async function executeQuery(query, params = []) {
  if (!pool) {
    throw new Error('数据库连接池未初始化');
  }

  try {
    const [rows, fields] = await pool.execute(query, params);
    return { rows, fields };
  } catch (error) {
    console.error('SQL执行错误:', error.message);
    throw error;
  }
}

/**
 * 获取数据库表列表
 * @returns {Promise<Array>} 表名列表
 */
async function getTables() {
  const result = await executeQuery('SHOW TABLES');
  return result.rows.map((row) => Object.values(row)[0]);
}

/**
 * 获取表结构信息
 * @param {string} tableName - 表名
 * @returns {Promise<Array>} 表结构信息
 */
async function getTableSchema(tableName) {
  const result = await executeQuery('DESCRIBE ??', [tableName]);
  return result.rows;
}

/**
 * 创建MCP服务器实例
 */
const server = new Server(
  {
    name: 'mcp-server-mysql',
    version: '9.4.0',
  },
  {
    capabilities: {
      tools: {},
    },
  }
);

/**
 * 注册工具列表处理器
 */
server.setRequestHandler(ListToolsRequestSchema, async () => {
  return {
    tools: [
      {
        name: 'mysql_query',
        description: '执行MySQL查询语句',
        inputSchema: {
          type: 'object',
          properties: {
            query: {
              type: 'string',
              description: 'SQL查询语句',
            },
            params: {
              type: 'array',
              description: '查询参数（可选）',
              items: {
                type: 'string',
              },
            },
          },
          required: ['query'],
        },
      },
      {
        name: 'mysql_tables',
        description: '获取数据库中所有表的列表',
        inputSchema: {
          type: 'object',
          properties: {},
        },
      },
      {
        name: 'mysql_schema',
        description: '获取指定表的结构信息',
        inputSchema: {
          type: 'object',
          properties: {
            table: {
              type: 'string',
              description: '表名',
            },
          },
          required: ['table'],
        },
      },
      {
        name: 'mysql_insert',
        description: '向表中插入数据',
        inputSchema: {
          type: 'object',
          properties: {
            table: {
              type: 'string',
              description: '表名',
            },
            data: {
              type: 'object',
              description: '要插入的数据（键值对）',
            },
          },
          required: ['table', 'data'],
        },
      },
      {
        name: 'mysql_update',
        description: '更新表中的数据',
        inputSchema: {
          type: 'object',
          properties: {
            table: {
              type: 'string',
              description: '表名',
            },
            data: {
              type: 'object',
              description: '要更新的数据（键值对）',
            },
            where: {
              type: 'object',
              description: '更新条件（键值对）',
            },
          },
          required: ['table', 'data', 'where'],
        },
      },
      {
        name: 'mysql_delete',
        description: '删除表中的数据',
        inputSchema: {
          type: 'object',
          properties: {
            table: {
              type: 'string',
              description: '表名',
            },
            where: {
              type: 'object',
              description: '删除条件（键值对）',
            },
          },
          required: ['table', 'where'],
        },
      },
    ],
  };
});

/**
 * 注册工具调用处理器
 */
server.setRequestHandler(CallToolRequestSchema, async (request) => {
  const { name, arguments: args } = request.params;

  try {
    switch (name) {
      case 'mysql_query': {
        const { query, params = [] } = args;
        const result = await executeQuery(query, params);
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(
                {
                  rows: result.rows,
                  rowCount: Array.isArray(result.rows)
                    ? result.rows.length
                    : result.rows.affectedRows || 0,
                  fields:
                    result.fields?.map((f) => ({
                      name: f.name,
                      type: f.type,
                    })) || [],
                },
                null,
                2
              ),
            },
          ],
        };
      }

      case 'mysql_tables': {
        const tables = await getTables();
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({ tables }, null, 2),
            },
          ],
        };
      }

      case 'mysql_schema': {
        const { table } = args;
        const schema = await getTableSchema(table);
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({ table, schema }, null, 2),
            },
          ],
        };
      }

      case 'mysql_insert': {
        const { table, data } = args;
        const columns = Object.keys(data);
        const values = Object.values(data);
        const placeholders = columns.map(() => '?').join(', ');

        const query = `INSERT INTO ?? (${columns
          .map(() => '??')
          .join(', ')}) VALUES (${placeholders})`;
        const params = [table, ...columns, ...values];

        const result = await executeQuery(query, params);
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(
                {
                  success: true,
                  insertId: result.rows.insertId,
                  affectedRows: result.rows.affectedRows,
                },
                null,
                2
              ),
            },
          ],
        };
      }

      case 'mysql_update': {
        const { table, data, where } = args;
        const setClause = Object.keys(data)
          .map(() => '?? = ?')
          .join(', ');
        const whereClause = Object.keys(where)
          .map(() => '?? = ?')
          .join(' AND ');

        const query = `UPDATE ?? SET ${setClause} WHERE ${whereClause}`;
        const setParams = Object.entries(data).flat();
        const whereParams = Object.entries(where).flat();
        const params = [table, ...setParams, ...whereParams];

        const result = await executeQuery(query, params);
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(
                {
                  success: true,
                  affectedRows: result.rows.affectedRows,
                  changedRows: result.rows.changedRows,
                },
                null,
                2
              ),
            },
          ],
        };
      }

      case 'mysql_delete': {
        const { table, where } = args;
        const whereClause = Object.keys(where)
          .map(() => '?? = ?')
          .join(' AND ');

        const query = `DELETE FROM ?? WHERE ${whereClause}`;
        const whereParams = Object.entries(where).flat();
        const params = [table, ...whereParams];

        const result = await executeQuery(query, params);
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(
                {
                  success: true,
                  affectedRows: result.rows.affectedRows,
                },
                null,
                2
              ),
            },
          ],
        };
      }

      default:
        throw new McpError(ErrorCode.MethodNotFound, `未知工具: ${name}`);
    }
  } catch (error) {
    throw new McpError(
      ErrorCode.InternalError,
      `工具执行错误: ${error.message}`
    );
  }
});

/**
 * 启动服务器
 */
async function main() {
  try {
    // 初始化数据库连接
    await initializeDatabase();

    // 创建传输层
    const transport = new StdioServerTransport();

    // 连接服务器和传输层
    await server.connect(transport);

    console.log('MCP MySQL服务器启动成功');
  } catch (error) {
    console.error('服务器启动失败:', error);
    process.exit(1);
  }
}

// 处理进程退出
process.on('SIGINT', async () => {
  console.log('\n正在关闭服务器...');
  if (pool) {
    await pool.end();
  }
  process.exit(0);
});

// 启动服务器
main().catch(console.error);
